import matplotlib.pyplot as plt
import matplotlib.image as img
import numpy as np
import cv2


# 大津阈值寻找
def OTSU_enhance(img_gray, th_begin=0, th_end=256, th_step=1):
    assert img_gray.ndim == 2, "must input a gary_img"  # 注意必须输入灰度图，对应前文中使用的rgb转灰度
    count = 0
    max_g = 0
    suitable_th = 0  # 自适应阈值
    # 遍历生成阈值，循环
    for threshold in range(th_begin, th_end, th_step):
        img_dx = img_gray > threshold
        img_bj = img_gray <= threshold
        # 求和
        fore_pix = np.sum(img_dx)
        back_pix = np.sum(img_bj)
        # 退出条件
        if 0 == fore_pix:
            break
        if 0 == back_pix:
            continue
        # （大津阈值算法）求最大方差
        w0 = float(fore_pix) / img_gray.size
        u0 = float(np.sum(img_gray * img_dx)) / fore_pix
        w1 = float(back_pix) / img_gray.size
        u1 = float(np.sum(img_gray * img_bj)) / back_pix
        # 内部类方差计算
        g = w0 * w1 * (u0 - u1) * (u0 - u1)
        if g > max_g:
            max_g = g
            suitable_th = threshold
    return suitable_th


def do_segment(img_gray, threshold):
    d = img_gray.copy()
    height, width = img_gray.shape
    for i in range(height):
        for j in range(width):
            if d[i, j] < threshold:
                d[i, j] = 0
            else:
                d[i, j] = 255
    return d


def whole_hist(image):
    """
	绘制整幅图像的直方图
	"""
    plt.hist(image.ravel(), 256, [0, 256])
    plt.show()


if __name__ == '__main__':
    img = cv2.imread('4.jpg')
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    d = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
    # 绘制直方图
    whole_hist(d)
    sp = d.shape
    print(sp)
    height = sp[0]
    weight = sp[1]
    threshold = OTSU_enhance(d, th_begin=0, th_end=256, th_step=1)  # 若要实现固定阈值分割，只需要把此处的阈值设为固定值，实验时以100为例
    print("用OTSU求出的阈值为", threshold)
    for i in range(height):
        for j in range(weight):
            if d[i, j] < threshold:
                d[i, j] = 0
            else:
                d[i, j] = 1
    plt.subplot(121), plt.imshow(img), plt.title('original')
    plt.subplot(122), plt.imshow(d, cmap=plt.cm.gray), plt.title('OTSU')
    plt.show()
